# -*- coding: utf-8 -*-
'''
作者：Zion Loh
版本：V1.1.1
License: MIT
使用该源代码请附带作者及版权信息！
'''

import asyncio
import asyncssh
from datetime import datetime
import os
import logging
import random
from tqdm.asyncio import tqdm_asyncio

# 初始化日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

async def restart_switch(hostname, username, password,
                         connect_timeout=10, command_timeout=120,
                         max_retries=3, retry_delay=(1, 3)):
    """
    异步重启交换机，带重试机制。
    max_retries: 最大重试次数
    retry_delay: 重试间隔秒数范围，随机等待，防止并发冲突
    """
    for attempt in range(1, max_retries + 1):
        try:
            logger.info(f"[{hostname}] 第{attempt}次尝试连接")
            async with asyncssh.connect(
                hostname,
                username=username,
                password=password,
                known_hosts=None,
                connect_timeout=connect_timeout
            ) as conn:
                logger.info(f"[{hostname}] 连接成功，发送重启命令")
                try:
                    cmd = "reboot"
                    result = await asyncio.wait_for(conn.run(cmd, check=True), timeout=command_timeout)
                    logger.info(f"[{hostname}] 重启命令执行成功")
                    return {'hostname': hostname, 'reboot_status': 'success'}
                except asyncio.TimeoutError:
                    logger.warning(f"[{hostname}] 重启命令执行超时")
                    return {'hostname': hostname, 'reboot_status': 'timeout', 'failure_reason': '命令执行超时'}
                except asyncssh.ProcessError as e:
                    logger.error(f"[{hostname}] 重启命令执行失败: {e}")
                    return {'hostname': hostname, 'reboot_status': 'failure', 'failure_reason': f"命令失败: {str(e)}"}

        except (asyncssh.Error, asyncio.TimeoutError) as e:
            logger.warning(f"[{hostname}] 连接失败: {e}")
            if attempt < max_retries:
                delay = random.uniform(*retry_delay)
                logger.info(f"[{hostname}] 第{attempt}次重试，等待{delay:.1f}秒")
                await asyncio.sleep(delay)
            else:
                logger.error(f"[{hostname}] 达到最大重试次数，放弃连接")
                return {'hostname': hostname, 'reboot_status': 'failure', 'failure_reason': f"连接失败: {str(e)}"}

async def monitor_switches(hostnames, username, password):
    tasks = [restart_switch(hostname, username, password) for hostname in hostnames]
    # 使用 tqdm_asyncio 展示异步任务进度条
    results = []
    for f in tqdm_asyncio.as_completed(tasks, total=len(tasks), desc="重启进度"):
        result = await f
        results.append(result)
    return results

def generate_html_report(results):
    today = datetime.now().strftime("%Y%m%d")
    report_filename = f"restart_report_{today}.html"

    if os.path.exists(report_filename):
        os.remove(report_filename)
        logger.info(f"删除旧的报告文件: {report_filename}")

    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    success_rows = []
    failure_rows = []

    for res in results:
        hostname = res.get('hostname', '未知主机')
        status = res.get('reboot_status', 'unknown')
        if status == 'success':
            success_rows.append(f"<tr><td>{hostname}</td><td>重启完成</td></tr>")
        else:
            reason = res.get('failure_reason', '未知原因')
            if status == 'timeout':
                failure_rows.append(f"<tr><td>{hostname}</td><td>重启执行超时</td><td>{reason}</td></tr>")
            else:
                failure_rows.append(f"<tr><td>{hostname}</td><td>重启失败</td><td>{reason}</td></tr>")

    html_content = f"""
<html>
<head><title>交换机批量重启报告</title></head>
<body>
<h2>交换机批量重启报告 - {timestamp}</h2>

<h3>成功重启的交换机：</h3>
<table border="1">
<tr><th>交换机名称</th><th>重启状态</th></tr>
{''.join(success_rows)}
</table>

<h3>重启失败的交换机：</h3>
<table border="1">
<tr><th>交换机名称</th><th>重启状态</th><th>失败原因</th></tr>
{''.join(failure_rows)}
</table>

</body>
</html>
"""

    with open(report_filename, "w", encoding="utf-8") as f:
        f.write(html_content)

    logger.info(f"HTML报告已生成：{report_filename}")

async def main():
    switch_hostnames = [
        '172.16.1.26', '172.16.1.27', '172.16.1.28', '172.16.1.29', '172.16.1.30',
        '172.16.1.3', '172.16.1.5', '172.16.1.6', '172.16.1.16', '172.16.1.17',
        '172.16.1.18', '172.16.1.19', '172.16.1.20', '172.16.1.21', '172.16.1.4',
        '172.16.1.7', '172.16.1.8', '172.16.1.9', '172.16.1.10', '172.16.1.12',
        '172.16.1.13', '172.16.1.14', '172.16.1.15', '172.16.1.11', '172.16.1.2',
        '172.16.1.1'
    ]
    username = 'admin'
    password = 'admin'

    logger.info("开始批量重启交换机...")
    results = await monitor_switches(switch_hostnames, username, password)
    logger.info("所有交换机重启完成！正在生成报告...")
    generate_html_report(results)

if __name__ == "__main__":
    asyncio.run(main())
